#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
使用Ollama本地Qwen3模型学习谱面创作规律
支持增量式学习方法
支持谱面段落识别和复用
"""

import os
import json
import glob
import re
import numpy as np
from typing import List, Dict, Any, Tuple
from openai import OpenAI
from collections import Counter

# 导入音频分析模块
from beat_analyzer import analyze_audio_features

class ChartSegment:
    """谱面段落类"""
    def __init__(self, notes: List[str], bpm: float = None, metadata: Dict = None):
        self.notes = notes
        self.bpm = bpm
        self.metadata = metadata or {}
        self.frequency = 1
        self.source_charts = []  # 来源谱面
        self.vector = None  # 向量表示
        self.normalized_notes = None  # 标准化后的音符
        
    def __hash__(self):
        # 基于标准化音符内容生成哈希值
        if self.normalized_notes:
            return hash(tuple(self.normalized_notes))
        return hash(tuple(self.notes))
        
    def __eq__(self, other):
        # 比较两个段落是否相同（基于标准化音符）
        if not isinstance(other, ChartSegment):
            return False
        if self.normalized_notes and other.normalized_notes:
            return self.normalized_notes == other.normalized_notes
        return self.notes == other.notes

class ChartLearningEngine:
    def __init__(self, data_dir: str = "data"):
        """
        初始化谱面学习引擎
        
        Args:
            data_dir: 包含谱面JSON文件的目录
        """
        self.data_dir = data_dir
        self.chart_files = []
        self.accumulated_knowledge = ""
        self.start_index = 0
        self.common_segments = {}  # 常见谱面段落 {hash: ChartSegment}
        self.segment_patterns = {}  # 段落模式统计
        
        # 初始化Ollama客户端
        self.client = OpenAI(
            base_url='http://localhost:11434/v1',
            api_key='ollama'  # Ollama不需要API密钥，但openai库需要一个值
        )
    
    def normalize_notes(self, notes: List[str]) -> List[str]:
        """
        标准化音符，消除位置平移影响
        
        Args:
            notes: 音符列表
            
        Returns:
            标准化后的音符列表
        """
        if not notes:
            return notes
            
        normalized = []
        # 解析所有音符，找出最小按键位置
        all_positions = []
        parsed_lines = []
        
        for line in notes:
            if line.startswith('{') and '},' in line:
                try:
                    # 提取note值
                    note_value_str = line[1:line.index('}')]
                    note_value = int(note_value_str)
                    # 提取音符部分
                    notes_part = line[line.index('},') + 2:]
                    note_list = notes_part.split(',') if notes_part else []
                    parsed_lines.append((note_value, note_list))
                    
                    # 收集所有按键位置
                    for note in note_list:
                        if note and note[0].isdigit():
                            pos = int(note[0])
                            all_positions.append(pos)
                except:
                    parsed_lines.append(line)
            else:
                parsed_lines.append(line)
        
        if not all_positions:
            return notes
            
        # 计算位置偏移（将最小位置移到1）
        min_pos = min(all_positions)
        offset = 1 - min_pos
        
        # 应用偏移
        for item in parsed_lines:
            if isinstance(item, tuple):
                note_value, note_list = item
                new_notes = []
                for note in note_list:
                    # 处理各种类型的音符
                    if note and note[0].isdigit():
                        pos = int(note[0])
                        new_pos = pos + offset
                        # 确保位置在1-8范围内
                        if 1 <= new_pos <= 8:
                            new_note = str(new_pos) + note[1:]
                        else:
                            # 如果超出范围，保持原样或做适当处理
                            new_note = note
                        new_notes.append(new_note)
                    else:
                        new_notes.append(note)
                
                # 重新构建行
                notes_str = ','.join(new_notes)
                normalized.append("{" + str(note_value) + "}," + notes_str)
            else:
                normalized.append(item)
                
        return normalized
    
    def vectorize_segment(self, notes: List[str]) -> np.ndarray:
        """
        将谱面段落转换为向量表示
        
        Args:
            notes: 音符列表
            
        Returns:
            向量表示
        """
        # 简化的向量化方法：统计各种音符类型的数量
        note_stats = {
            'tap': 0,
            'hold': 0,
            'slide': 0,
            'break': 0,
            'multi': 0,
            'total': len(notes)
        }
        
        for line in notes:
            if line.startswith('{') and '},' in line:
                try:
                    notes_part = line[line.index('},') + 2:]
                    note_list = notes_part.split(',') if notes_part else []
                    for note in note_list:
                        if not note:
                            continue
                        if 'h[' in note:
                            note_stats['hold'] += 1
                        elif 's' in note or '>' in note or '<' in note or '-' in note:
                            note_stats['slide'] += 1
                        elif 'b' in note:
                            note_stats['break'] += 1
                        elif '/' in note:
                            note_stats['multi'] += 1
                        elif note[0].isdigit():
                            note_stats['tap'] += 1
                except:
                    pass
        
        # 转换为向量
        vector = np.array([
            note_stats['tap'],
            note_stats['hold'],
            note_stats['slide'],
            note_stats['break'],
            note_stats['multi'],
            note_stats['total']
        ], dtype=float)
        
        # 归一化
        if note_stats['total'] > 0:
            vector[:-1] /= note_stats['total']
            
        return vector
    
    def segments_similar(self, seg1: ChartSegment, seg2: ChartSegment, threshold: float = 0.9) -> bool:
        """
        判断两个谱面段落是否相似
        
        Args:
            seg1: 第一个段落
            seg2: 第二个段落
            threshold: 相似度阈值
            
        Returns:
            是否相似
        """
        # 首先检查标准化音符
        if seg1.normalized_notes and seg2.normalized_notes:
            if seg1.normalized_notes == seg2.normalized_notes:
                return True
        
        # 检查向量相似度
        if seg1.vector is not None and seg2.vector is not None:
            # 计算余弦相似度
            dot_product = np.dot(seg1.vector, seg2.vector)
            norm1 = np.linalg.norm(seg1.vector)
            norm2 = np.linalg.norm(seg2.vector)
            
            if norm1 > 0 and norm2 > 0:
                cosine_sim = dot_product / (norm1 * norm2)
                return cosine_sim >= threshold
                
        return False
    
    def load_progress(self) -> bool:
        """
        从进度文件加载学习进度
        
        Returns:
            是否成功加载进度
        """
        progress_file = "learning_progress_ollama.json"
        if not os.path.exists(progress_file):
            print("未找到进度文件，将从头开始学习")
            return False
        
        try:
            with open(progress_file, 'r', encoding='utf-8') as f:
                progress_data = json.load(f)
            
            self.accumulated_knowledge = progress_data.get("accumulated_knowledge", "")
            self.start_index = progress_data.get("current_index", 0)
            # 加载常见段落
            if "common_segments" in progress_data:
                self.common_segments = {}
                for seg_hash, seg_data in progress_data["common_segments"].items():
                    segment = ChartSegment(seg_data["notes"], seg_data["bpm"], seg_data["metadata"])
                    segment.frequency = seg_data["frequency"]
                    segment.source_charts = seg_data["source_charts"]
                    segment.normalized_notes = seg_data.get("normalized_notes")
                    vector_data = seg_data.get("vector")
                    if vector_data:
                        segment.vector = np.array(vector_data)
                    self.common_segments[seg_hash] = segment
            
            print(f"已从进度文件加载学习进度: {self.start_index}")
            if self.accumulated_knowledge:
                print(f"已加载的知识摘要: {self.accumulated_knowledge[:200]}...")
            return True
        except Exception as e:
            print(f"加载进度文件时出错: {e}")
            return False
    
    def scan_chart_files(self):
        """扫描data目录下的所有谱面JSON文件"""
        pattern = os.path.join(self.data_dir, "*.json")
        self.chart_files = glob.glob(pattern)
        self.chart_files.sort()  # 按文件名排序
        
        print(f"总共找到 {len(self.chart_files)} 个谱面文件")
        return self.chart_files
    
    def read_chart_file(self, chart_file: str) -> Dict[str, Any]:
        """
        读取谱面JSON文件内容
        
        Args:
            chart_file: 谱面JSON文件路径
            
        Returns:
            谱面文件内容（JSON对象）
        """
        try:
            with open(chart_file, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception as e:
            print(f"读取谱面文件 {chart_file} 时出错: {e}")
            return {}
    
    def extract_chart_content(self, chart_data: Dict[str, Any]) -> str:
        """
        从JSON数据中提取level值最大的谱面内容
        
        Args:
            chart_data: 谱面JSON数据
            
        Returns:
            谱面内容字符串
        """
        # 在charts数组中查找level值最大的谱面
        charts = chart_data.get("chart_info", {}).get("charts", [])
        if not charts:
            return ""
        
        # 找到level值最大的谱面
        # 注意：level可能是字符串，需要处理类似"12"和"12+"这样的值
        def parse_level(chart):
            level_str = chart.get("level", "0")
            # 移除"+"号并转换为浮点数进行比较
            level_val = float(level_str.replace("+", ""))
            return level_val
            
        max_level_chart = max(charts, key=parse_level)
        chart_content = max_level_chart.get("chart", "")
        
        # 如果提取到的谱面内容为空，打印调试信息
        if not chart_content:
            print(f"警告: 未能提取到谱面内容，charts数量: {len(charts)}")
            for i, chart in enumerate(charts):
                print(f"  谱面 {i}: difficulty={chart.get('difficulty')}, level={chart.get('level')}")
        
        return chart_content
    
    def extract_chart_segments(self, chart_content: str, min_segment_length: int = 8, max_segment_length: int = 32) -> List[ChartSegment]:
        """
        从谱面内容中提取段落
        
        Args:
            chart_content: 谱面内容
            min_segment_length: 最小段落长度
            max_segment_length: 最大段落长度
            
        Returns:
            谱面段落列表
        """
        segments = []
        lines = chart_content.split('\n')
        
        # 提取谱面行（去除元数据和结束标记）
        chart_lines = []
        bpm = None
        
        for line in lines:
            line = line.strip()
            if line.startswith('&wholebpm='):
                try:
                    bpm = float(line.split('=')[1])
                except:
                    bpm = 120.0
            elif line.startswith('{') and '},' in line:
                chart_lines.append(line)
            # 添加对没有&wholebpm前缀但包含节拍信息的行的处理
            elif line.startswith('{') and ',' in line:
                chart_lines.append(line)
        
        print(f"从谱面中提取到 {len(chart_lines)} 行有效谱面数据")
        
        # 生成不同长度的段落
        for length in range(min_segment_length, min(max_segment_length + 1, len(chart_lines) + 1)):
            for i in range(len(chart_lines) - length + 1):
                segment_lines = chart_lines[i:i + length]
                segment = ChartSegment(segment_lines, bpm)
                # 添加标准化音符和向量表示
                segment.normalized_notes = self.normalize_notes(segment_lines)
                segment.vector = self.vectorize_segment(segment_lines)
                segments.append(segment)
                
        return segments
    
    def identify_common_segments(self, new_segments: List[ChartSegment], chart_name: str):
        """
        识别常见谱面段落
        
        Args:
            new_segments: 新的谱面段落
            chart_name: 谱面名称
        """
        for segment in new_segments:
            seg_hash = hash(segment)
            found_similar = False
            
            # 首先检查完全相同的段落
            if seg_hash in self.common_segments:
                # 段落已存在，增加频率
                self.common_segments[seg_hash].frequency += 1
                if chart_name not in self.common_segments[seg_hash].source_charts:
                    self.common_segments[seg_hash].source_charts.append(chart_name)
                found_similar = True
            else:
                # 检查相似的段落
                for existing_hash, existing_segment in self.common_segments.items():
                    if self.segments_similar(segment, existing_segment):
                        # 找到相似段落，增加频率
                        existing_segment.frequency += 1
                        if chart_name not in existing_segment.source_charts:
                            existing_segment.source_charts.append(chart_name)
                        found_similar = True
                        break
            
            # 如果没有找到相似段落，则添加新段落
            if not found_similar:
                segment.source_charts.append(chart_name)
                self.common_segments[seg_hash] = segment
    
    def analyze_audio_for_chart(self, chart_data: Dict[str, Any]) -> Dict[str, Any]:
        """
        从谱面数据中获取音频特征
        
        Args:
            chart_data: 谱面JSON数据
            
        Returns:
            音频特征数据
        """
        # 直接从JSON数据中获取音频特征
        audio_features = chart_data.get("audio_features", {})
        if audio_features:
            print(f"已从JSON数据中获取音频特征")
            return audio_features
        
        # 如果JSON中没有音频特征，则尝试从音频文件分析
        audio_file = chart_data.get("audio_file", "")
        if audio_file and os.path.exists(audio_file):
            try:
                features = analyze_audio_features(audio_file)
                print(f"已通过音频文件分析获取特征: {audio_file}")
                return features
            except Exception as e:
                print(f"分析音频文件 {audio_file} 时出错: {e}")
        
        # 返回空特征
        print("无法获取音频特征")
        return {}
    
    def send_to_ollama(self, prompt: str) -> str:
        """
        发送请求到Ollama本地模型
        
        Args:
            prompt: 发送给模型的提示
            
        Returns:
            模型响应
        """
        try:
            response = self.client.chat.completions.create(
                model="qwen3:1.7b",
                messages=[
                    {"role": "system", "content": "你是一个专业的音乐游戏谱面分析师，能够深入分析谱面结构和创作规律。"},
                    {"role": "user", "content": prompt}
                ],
                temperature=0.7,
                max_tokens=2000
            )
            return response.choices[0].message.content
        except Exception as e:
            print(f"调用Ollama模型时出错: {e}")
            return ""
    
    def save_progress(self, current_index: int):
        """
        保存学习进度到文件
        
        Args:
            current_index: 当前处理到的文件索引
        """
        progress_data = {
            "accumulated_knowledge": self.accumulated_knowledge,
            "current_index": current_index,
            "common_segments": {}
        }
        
        # 保存常见段落
        for seg_hash, segment in self.common_segments.items():
            progress_data["common_segments"][str(seg_hash)] = {
                "notes": segment.notes,
                "bpm": segment.bpm,
                "metadata": segment.metadata,
                "frequency": segment.frequency,
                "source_charts": segment.source_charts,
                "normalized_notes": segment.normalized_notes,
                "vector": segment.vector.tolist() if segment.vector is not None else None
            }
        
        progress_file = "learning_progress_ollama.json"
        try:
            with open(progress_file, 'w', encoding='utf-8') as f:
                json.dump(progress_data, f, ensure_ascii=False, indent=2)
            print(f"进度已保存到 {progress_file}")
        except Exception as e:
            print(f"保存进度文件时出错: {e}")
    
    def learn_chart_patterns(self, chart_file: str) -> str:
        """
        学习单个谱面的创作规律
        
        Args:
            chart_file: 谱面文件路径
            
        Returns:
            学习到的规律知识
        """
        print(f"正在学习谱面: {chart_file}")
        
        # 读取谱面文件
        chart_data = self.read_chart_file(chart_file)
        if not chart_data:
            return ""
        
        # 获取音频特征
        audio_features = self.analyze_audio_for_chart(chart_data)
        
        # 提取谱面内容（只提取level值最大的谱面）
        chart_content = self.extract_chart_content(chart_data)
        if not chart_content:
            print(f"谱面文件 {chart_file} 中未找到有效的谱面内容")
            return ""
        
        # 提取谱面段落
        segments = self.extract_chart_segments(chart_content)
        print(f"提取到 {len(segments)} 个谱面段落")
        
        # 识别常见段落
        self.identify_common_segments(segments, os.path.basename(chart_file))
        
        # 构造提示词
        prompt = f"请分析以下音乐游戏谱面的创作规律和特点：\n\n"
        prompt += f"谱面文件: {os.path.basename(chart_file)}\n\n"
        prompt += f"音频特征:\n{json.dumps(audio_features, ensure_ascii=False, indent=2)}\n\n"
        prompt += f"谱面内容:\n{chart_content[:2000]}\n\n"  # 限制长度避免过长
        
        prompt += "请从以下几个方面进行分析：\n"
        prompt += "1. 谱面整体结构和段落划分\n"
        prompt += "2. 音符配置模式和常用技巧\n"
        prompt += "3. 谱面密度变化规律\n"
        prompt += "4. 与音频特征的对应关系\n"
        prompt += "5. 可能的创作思路和设计意图\n\n"
        prompt += "请用简洁明了的语言总结，便于后续创作参考。"
        
        # 发送到Ollama模型
        response = self.send_to_ollama(prompt)
        if response:
            print(f"已完成谱面 {chart_file} 的学习和知识整合")
            return response
        else:
            print(f"谱面 {chart_file} 学习过程中出现错误")
            return ""
    
    def find_top_segments(self, top_n: int = 10) -> List[Tuple[int, ChartSegment]]:
        """
        找出出现频率最高的段落
        
        Args:
            top_n: 返回前N个高频段落
            
        Returns:
            [(频率, 段落对象), ...] 按频率降序排列
        """
        # 按频率排序
        sorted_segments = sorted(
            self.common_segments.items(),
            key=lambda x: x[1].frequency,
            reverse=True
        )
        return [(seg.frequency, seg) for _, seg in sorted_segments[:top_n]]
    
    def run_learning(self, batch_size: int = 50):
        """
        运行谱面学习过程
        
        Args:
            batch_size: 每批处理的谱面数量
        """
        print("开始使用Ollama本地Qwen3模型进行增量式谱面学习")
        print("=" * 50)
        
        # 扫描谱面文件
        self.scan_chart_files()
        
        if not self.chart_files:
            print("未找到任何谱面文件")
            return
        
        # 加载进度
        self.load_progress()
        
        # 确定起始位置
        start_idx = max(0, self.start_index)
        total_files = len(self.chart_files)
        
        print(f"将处理 {total_files} 个谱面")
        
        # 分批处理谱面文件
        for i in range(start_idx, total_files, batch_size):
            batch_end = min(i + batch_size, total_files)
            print(f"\n处理进度: {i+1}-{batch_end}/{total_files}")
            
            # 处理当前批次
            batch_knowledge = []
            for j in range(i, batch_end):
                chart_file = self.chart_files[j]
                print(f"  正在学习谱面: {chart_file}")
                
                knowledge = self.learn_chart_patterns(chart_file)
                if knowledge:
                    batch_knowledge.append(knowledge)
                
                # 每处理一个文件就保存进度
                self.save_progress(j + 1)
            
            # 整合批次知识
            if batch_knowledge:
                batch_prompt = "请整合以下" + str(len(batch_knowledge)) + "个谱面的学习成果，总结通用的创作规律：\n\n"
                batch_prompt += "学习成果:\n" + '\n---\n'.join(batch_knowledge) + "\n\n"
                batch_prompt += "请总结这些谱面的共同特点和创作规律，包括但不限于:\n"
                batch_prompt += "1. 常见的谱面结构模式\n"
                batch_prompt += "2. 高频出现的音符配置技巧\n"
                batch_prompt += "3. 谱面密度和难度分布规律\n"
                batch_prompt += "4. 与音乐特征的配合方式\n"
                batch_prompt += "5. 创作者常用的设计思路\n"
                
                batch_summary = self.send_to_ollama(batch_prompt)
                if batch_summary:
                    self.accumulated_knowledge += f"\n\n批次 {i//batch_size + 1} 总结:\n{batch_summary}"
        
        # 保存最终进度
        self.save_progress(total_files)
        
        # 显示高频段落
        print("\n" + "=" * 50)
        print("学习完成！发现的高频谱面段落:")
        top_segments = self.find_top_segments(10)
        for i, (freq, segment) in enumerate(top_segments, 1):
            print(f"{i}. 出现 {freq} 次: {segment.notes[:3]}...")
        
        # 保存最终报告
        self.generate_final_report()
    
    def generate_final_report(self):
        """生成最终学习报告"""
        report = {
            "accumulated_knowledge": self.accumulated_knowledge,
            "total_segments": len(self.common_segments),
            "top_patterns": []
        }
        
        # 添加高频段落
        top_segments = self.find_top_segments(20)
        for freq, segment in top_segments:
            report["top_patterns"].append({
                "frequency": freq,
                "notes": segment.notes,
                "bpm": segment.bpm,
                "source_charts": segment.source_charts
            })
        
        # 保存报告
        with open("learning_report_ollama.json", "w", encoding="utf-8") as f:
            json.dump(report, f, ensure_ascii=False, indent=2)
        
        print("最终学习报告已保存到 learning_report_ollama.json")

def main():
    print("启动脚本...")
    
    # 创建学习引擎实例
    engine = ChartLearningEngine("data")
    
    # 运行学习过程，处理所有谱面文件
    engine.run_learning(batch_size=50)  # 增大批处理大小以提高效率

if __name__ == "__main__":
    main()